import  torch
from    torch import optim, nn
# import  visdom
import  torchvision
from    torch.utils.data import DataLoader
from datam import *
from    resnet import ResNet18
from config import config,getData,getPolemon
import wandb
import yaml


data_dir = config.data_dir
batchsz = config.BATCH_SIZE
lr = config.lr
epochs = config.max_epochs
device = torch.device('cuda')
torch.manual_seed(1234)
train_loader,val_loader,test_loader = getPolemon()
# train_loader,val_loader,test_loader = getData()

def evalute(model, loader):
    model.eval()
    
    correct = 0
    total = len(loader.dataset)

    for x,y in loader:
        x,y = x.to(device), y.to(device)
        with torch.no_grad():
            logits = model(x)
            pred = logits.argmax(dim=1)
        correct += torch.eq(pred, y).sum().float().item()

    return correct / total

def main():

    model = ResNet18(5).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criteon = nn.CrossEntropyLoss()


    best_acc, best_epoch = 0, 0
    global_step = 0
    # viz.line([0], [-1], win='loss', opts=dict(title='loss'))
    # viz.line([0], [-1], win='val_acc', opts=dict(title='val_acc'))
    for epoch in range(epochs):

        for step, (x,y) in enumerate(train_loader):

            # x: [b, 3, 224, 224], y: [b]
            x, y = x.to(device), y.to(device)
            
            model.train()
            logits = model(x)
            loss = criteon(logits, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # viz.line([loss.item()], [global_step], win='loss', update='append')
            wandb.log({
                'epoch': epoch,
                'lr': lr,
                'loss': loss.item(),
            })
            global_step += 1

        if epoch % 1 == 0:

            val_acc = evalute(model, val_loader)
            if val_acc> best_acc:
                best_epoch = epoch
                best_acc = val_acc

                torch.save(model.state_dict(), 'best.mdl')

                # viz.line([val_acc], [global_step], win='val_acc', update='append')
                wandb.log({
                    'val_acc': val_acc,
                })


    print('best acc:', best_acc, 'best epoch:', best_epoch)

    model.load_state_dict(torch.load('best.mdl'))
    print('loaded from ckpt!')

    test_acc = evalute(model, test_loader)
    print('test acc:', test_acc)

if __name__ == '__main__':
    wandb.init(project=config.project_name)
    main()
    wandb.finish()